package javaagent.mybatis;

import javassist.*;
import org.apache.ibatis.mapping.BoundSql;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.security.ProtectionDomain;
import java.util.Date;

/**
 * @author ：zzs
 * @version : 1.0
 * @date ：Created in 2022/3/4 20:04
 * @description：
 */
public class MyBatisAgent implements ClassFileTransformer {

    @Override
    public byte[] transform(ClassLoader loader, String className,
                            Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) throws IllegalClassFormatException {
        if (!"org/apache/ibatis/executor/BaseExecutor".equals(className)) {
            return null;
        }
        System.out.println("进入agent功能");
        try {
            // 将jar包注入到与spring同一个类加载器中然后用tomcat进行启动
            // tomcat 中可行，在Spring boot中不可行,因为tomcat使用反射进行类之间访问,反射不受双亲委派机制限制
            appendToLoader(loader);
        } catch (Exception e) {
            System.err.println("jar 注入失败");
            e.printStackTrace();
            return null;
        }
        try {
            ClassPool pool = new ClassPool();
            pool.appendSystemPath();
            pool.appendClassPath(new LoaderClassPath(loader));
            CtClass ctClass = pool.get("org.apache.ibatis.executor.BaseExecutor");
            // 查询
            CtMethod ctMethod = ctClass.getDeclaredMethods("query")[1];
            ctMethod.addLocalVariable("info", pool.get(SqlInfo.class.getName()));
            ctMethod.insertBefore("info=javaagent.mybatis.MyBatisAgent.begin($args);");
            ctMethod.insertAfter("javaagent.mybatis.MyBatisAgent.end(info);");
            System.out.println("插桩成功：" + ctClass.getName());
            return ctClass.toBytecode();
        } catch (NotFoundException | CannotCompileException | IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    // 将agent.jar加入到与spring同一个加载器中
    private void appendToLoader(ClassLoader loader) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException, MalformedURLException {
        URLClassLoader urlClassLoader = (URLClassLoader) loader;
        Method addURL = URLClassLoader.class.getDeclaredMethod("addURL", URL.class);
        addURL.setAccessible(true);
        String path = MyBatisAgent.class.getResource("").getPath();
        path = path.substring(0, path.indexOf("!/"));
        addURL.invoke(urlClassLoader, new URL(path));
    }

    public static SqlInfo begin(Object[] params) {
        SqlInfo info = new SqlInfo();
        info.beginTime = System.currentTimeMillis();
        BoundSql adapter = (BoundSql) params[5];
        info.sql = adapter.getSql();
        return info;
    }

    public static void end(SqlInfo info) {
        info.useTime = System.currentTimeMillis() - info.beginTime;
        System.out.println(info);
    }

    public static class SqlInfo {
        public long beginTime;
        public long useTime;
        public String sql;

        @Override
        public String toString() {
            return "SqlInfo{" +
                    "beginTime=" + new Date(beginTime) +
                    ", useTime=" + useTime +
                    ", sql='" + sql + '\'' +
                    '}';
        }
    }
}
